// Copyright 2022 Chen Jun
// Licensed under the MIT License.

#ifndef ARMOR_DETECTOR_PNP_SOLVER_HPP_
#define ARMOR_DETECTOR_PNP_SOLVER_HPP_

#include "hnurm_armor/Compensator.h"
#include "hnurm_armor/DataType.hpp"
#include "hnurm_armor/armor.hpp"
#include "hnurm_armor/tracker.hpp"

#include <opencv2/calib3d.hpp>
#include <opencv2/core.hpp>
#include <opencv2/core/eigen.hpp>
#include <Eigen/Core>
#include <Eigen/Dense>

#include <angles/angles.h>

#include <cmath>
#include <cstdlib>
#include <geometry_msgs/msg/transform_stamped.hpp>
#include <tf2/LinearMath/Quaternion.h>
#include <tf2_eigen/tf2_eigen.hpp>
#include <vector>
#include <filesystem>   
#include <ament_index_cpp/get_package_share_directory.hpp>

namespace hnurm
{
    class TSolver
    {
    public:
        explicit TSolver(const rclcpp::Node::SharedPtr &node)
        {
            small_armor_points_ = {
                cv::Point3f(0, 0, SMALL_ARMOR_HEIGHT / 2),
                cv::Point3f(-SMALL_ARMOR_WIDTH / 2, 0, 0),
                cv::Point3f(0, 0, -SMALL_ARMOR_HEIGHT / 2),
                cv::Point3f(SMALL_ARMOR_WIDTH / 2, 0, 0),
                cv::Point3f(0, 0, 0),
            };
            small_armor_9points_ = {
                cv::Point3f(-SMALL_ARMOR_WIDTH / 2, 0, SMALL_ARMOR_HEIGHT / 2),
                cv::Point3f(-SMALL_ARMOR_WIDTH / 2, 0, -SMALL_ARMOR_HEIGHT / 2),
                cv::Point3f(SMALL_ARMOR_WIDTH / 2, 0, -SMALL_ARMOR_HEIGHT / 2),
                cv::Point3f(SMALL_ARMOR_WIDTH / 2, 0, SMALL_ARMOR_HEIGHT / 2),
                cv::Point3f(0, 0, SMALL_ARMOR_HEIGHT / 2),
                cv::Point3f(-SMALL_ARMOR_WIDTH / 2, 0, 0),
                cv::Point3f(0, 0, -SMALL_ARMOR_HEIGHT / 2),
                cv::Point3f(SMALL_ARMOR_WIDTH / 2, 0, 0),
                cv::Point3f(0, 0, 0),
            };
            large_armor_points_ = {
                cv::Point3f(0, 0, LARGE_ARMOR_WIDTH / 2),
                cv::Point3f(-LARGE_ARMOR_WIDTH / 2, 0, 0),
                cv::Point3f(0, 0, -LARGE_ARMOR_WIDTH / 2),
                cv::Point3f(LARGE_ARMOR_WIDTH / 2, 0, 0),
                cv::Point3f(0, 0, 0),
            };
            large_armor_9points_ = {
                cv::Point3f(-LARGE_ARMOR_WIDTH / 2, 0, LARGE_ARMOR_HEIGHT / 2),
                cv::Point3f(-LARGE_ARMOR_WIDTH / 2, 0, -LARGE_ARMOR_HEIGHT / 2),
                cv::Point3f(LARGE_ARMOR_WIDTH / 2, 0, -LARGE_ARMOR_HEIGHT / 2),
                cv::Point3f(LARGE_ARMOR_WIDTH / 2, 0, LARGE_ARMOR_HEIGHT / 2),
                cv::Point3f(0, 0, LARGE_ARMOR_WIDTH / 2),
                cv::Point3f(-LARGE_ARMOR_WIDTH / 2, 0, 0),
                cv::Point3f(0, 0, -LARGE_ARMOR_WIDTH / 2),
                cv::Point3f(LARGE_ARMOR_WIDTH / 2, 0, 0),
                cv::Point3f(0, 0, 0),
            };
            float cam_bias_z, cam_bias_y, cam_bias_x;
            // cam_bias_z = node->declare_parameter("processor.solver.camera_bias_z", 0);
            // cam_bias_y = node->declare_parameter("processor.solver.camera_bias_y", 0);
            // cam_bias_x = node->declare_parameter("processor.solver.camera_bias_x", 0);
            // _cam_bias << cam_bias_z, cam_bias_y, cam_bias_x;
            std::string base_dir = ament_index_cpp::get_package_share_directory("hnurm_armor");
            std::string camera_id = node->declare_parameter("processor.solver.camera_id", "0");
            base_dir += "/data/";
            std::filesystem::path calib_path(base_dir);
            if (!std::filesystem::exists(calib_path)) {
                RCLCPP_WARN_STREAM(node->get_logger(), "Base calibration directory does not exist: " << calib_path.string());
            }

            calib_path /= (camera_id + ".yaml");
            std::string calib_info_path = calib_path.string();

            cv::FileStorage calib_info(calib_info_path, cv::FileStorage::READ);
            if (!calib_info.isOpened()) {
                RCLCPP_ERROR_STREAM(node->get_logger(), "Failed to open calibration file: " << calib_info_path);
            }
            cv::Mat camera_matrix, dist_coeffs;
            calib_info["camera_matrix"] >> camera_matrix_;   
            calib_info["distortion_coefficients"] >> dist_coeffs_;

            // 检查数据是否有效
            if (camera_matrix_.empty() || dist_coeffs_.empty()) {
                RCLCPP_ERROR(node->get_logger(), "Failed to read calibration parameters from file");
            }

        RCLCPP_INFO_STREAM(node->get_logger(), "Loaded camera matrix:\n" << camera_matrix_);
        RCLCPP_INFO_STREAM(node->get_logger(), "Loaded distortion coefficients:\n" << dist_coeffs_);


    }
    float distance(const cv::Point2f &p1, const cv::Point2f &p2)
    {
        float dx = p1.x - p2.x;
        float dy = p1.y - p2.y;
        return sqrt(dx * dx + dy * dy);
    }

    float totalDistance(const std::vector<cv::Point2f> &points1, const std::vector<cv::Point2f> &points2)
    {
        float total_dist = 0.0f;
        for (size_t i = 0; i < points1.size(); ++i)
        {
            total_dist += distance(points1[i], points2[i]);
        }
        return total_dist;
    }
    static void rotationMatrixToEulerAngles(const Eigen::Matrix3d &R, Eigen::Vector3d &eulerAngles)
    {
        double sy = sqrt(R(0, 0) * R(0, 0) + R(1, 0) * R(1, 0));
        bool singular = sy < 1e-6; // If
        if (!singular)
        {
            eulerAngles[0] = atan2(R(2, 1), R(2, 2)); // pitch
            eulerAngles[1] = atan2(-R(2, 0), sy); // roll
            eulerAngles[2] = atan2(R(1, 0), R(0, 0)); // yaw
        }
        else
        {
            eulerAngles[0] = atan2(-R(1, 2), R(1, 1));
            eulerAngles[1] = atan2(-R(2, 0), sy);
            eulerAngles[2] = 0;
        }
    }

    float get_distance(const float &yaw_reg, Armor &armor, cv::Mat &tvec)
    {
        double pitch_deg;
        if (armor.number == "outpost")
        {
            pitch_deg = 15.0;
        }
        else
        {
            pitch_deg = -15.0;
        }
        double pitch_rad = pitch_deg * CV_PI / 180.0;
        double yaw_rad = yaw_reg * CV_PI / 180.0;
        Eigen::Matrix3d rotation_matrix_pitch_roll;
        rotation_matrix_pitch_roll << 1, 0, 0,
            0, cos(pitch_rad), -sin(pitch_rad),
            0, sin(pitch_rad), cos(pitch_rad);
        Eigen::Matrix3d rotation_matrix_yaw;
        rotation_matrix_yaw << cos(yaw_rad), -sin(yaw_rad), 0,
            sin(yaw_rad), cos(yaw_rad), 0,
            0, 0, 1;
        Eigen::Matrix3d final_rotation_matrix = rotation_matrix_yaw * rotation_matrix_pitch_roll;
        std::vector<cv::Point2f> image_points;
        cv::Matx33d final_rotation_matrix_opencv(
            final_rotation_matrix(0, 0), final_rotation_matrix(0, 1), final_rotation_matrix(0, 2),
            final_rotation_matrix(1, 0), final_rotation_matrix(1, 1), final_rotation_matrix(1, 2),
            final_rotation_matrix(2, 0), final_rotation_matrix(2, 1), final_rotation_matrix(2, 2)
        );
        cv::Matx33d r;
        cv::eigen2cv(armor._R, r);
        if (armor.armor_type == ArmorType::SMALL)
        {
            projectPoints(small_armor_9points_,
                          r * final_rotation_matrix_opencv,
                          tvec,
                          camera_matrix_,
                          dist_coeffs_,
                          image_points);
        }
        else
        {
            projectPoints(large_armor_9points_,
                          r * final_rotation_matrix_opencv,
                          tvec,
                          camera_matrix_,
                          dist_coeffs_,
                          image_points);
        }
//        std::vector<cv::Point2f> last_five_points(armor.points2d.end() - 5, armor.points2d.end());
        float distance_9 = totalDistance(armor.points2d, image_points);
        return distance_9;
    }
    // Get 3d position from 2d points
    bool GetTransformation(Armor &armor, float g_pitch, float g_yaw)
    {
        // 求解PnP问题
        //            g_pitch = 0;
        //                    g_yaw=0;
        cv::Mat rvec, tvec;
        if(armor.armor_type == ArmorType::SMALL)
        {
            cv::solvePnP(
                small_armor_9points_, armor.points2d, camera_matrix_, dist_coeffs_, rvec, tvec, false, cv::SOLVEPNP_IPPE
            );
//            solvePnPBA(small_armor_points_, armor.points2d, camera_matrix_, dist_coeffs_, rvec, tvec, true);
        }
        else if(armor.armor_type == ArmorType::LARGE)
            cv::solvePnP(
                large_armor_9points_, armor.points2d, camera_matrix_, dist_coeffs_, rvec, tvec, false, cv::SOLVEPNP_IPPE
            );

        // 构造变换矩阵
        armor._translation << tvec.at<double>(0, 0), tvec.at<double>(1, 0), tvec.at<double>(2, 0);
        Eigen::AngleAxisd rot_angle(
            cv::norm(rvec), Eigen::Vector3d(rvec.at<double>(0, 0), rvec.at<double>(1, 0), rvec.at<double>(2, 0))
        );
        armor._rmat = rot_angle.matrix();
        float sp = std::sin(g_pitch), cp = std::cos(g_pitch), sy = std::sin(g_yaw), cy = std::cos(g_yaw);
        armor._r1 << 1, 0, 0,
                    0, 0, 1,
                    0, -1, 0;
        armor._r2 << 1, 0, 0,
                    0, cp, -sp,
                    0, sp, cp;
        armor._r3 << cy, -sy, 0,
                    sy, cy, 0,
                    0, 0, 1;
        armor._R = (armor._r3 * armor._r2 * armor._r1).inverse();
        armor.rotation = armor._r3 * armor._r2 * armor._r1 * armor._rmat;
        armor.position = armor._r3 * armor._r2 * (armor._r1 * armor._translation) / 1'000;
        // // armor.position = armor._r3 * armor._r2 * (armor._r1 * armor._translation) / 1'000;
        // std::cout<<"bias:"<<std::endl;
        // std::cout<< _cam_bias / 1'000<<std::endl;
        // std::cout<<"posstion"<<std::endl;
        // std::cout<<armor._r3 * armor._r2 * (armor._r1 * armor._translation) / 1'000<<std::endl;
        if (armor.number != "base")
        {
            Eigen::Vector3d eulerAngles;
            rotationMatrixToEulerAngles(armor.rotation, eulerAngles);
            float a = eulerAngles[2] * 180 / CV_PI - 35;  // 搜索区间起点
            float b = eulerAngles[2] * 180 / CV_PI + 35;   // 搜索区间终点
            float tolerance = 0.1; // 精度阈值
            const float phi = (1 + std::sqrt(5)) / 2; // 黄金分割比
            float c = b - (b - a) / phi;
            float d = a + (b - a) / phi;
            while (std::fabs(b - a) > tolerance)
            {
                float s1 = get_distance(c, armor, tvec);
                float s2 = get_distance(d, armor, tvec);
                if (s1 < s2)
                {
                    b = d;
                }
                else
                {
                    a = c;
                }
                c = b - (b - a) / phi;
                d = a + (b - a) / phi;
            }
            float right_yaw = (a + b) / 2 * CV_PI / 180;
            Eigen::Matrix3d rotation_matrix_pitch_roll;
            double pitch_deg;
            if (armor.number == "outpost")
            {
                pitch_deg = 15.0;
            }
            else if (armor.number == "base")
            {
                pitch_deg = -62.5;
            }
            else
            {
                pitch_deg = -15.0;
            }
            rotation_matrix_pitch_roll << 1, 0, 0,
                0, cos(pitch_deg * CV_PI / 180.0), -sin(pitch_deg * CV_PI / 180.0),
                0, sin(pitch_deg * CV_PI / 180.0), cos(pitch_deg * CV_PI / 180.0);
            Eigen::Matrix3d rotation_matrix_yaw;
            rotation_matrix_yaw << cos(right_yaw), -sin(right_yaw), 0,
                sin(right_yaw), cos(right_yaw), 0,
                0, 0, 1;
            armor.rotation = rotation_matrix_yaw * rotation_matrix_pitch_roll;
            return true;
        }
        return true;
    }


private:
    cv::Mat camera_matrix_;
    cv::Mat dist_coeffs_;

    static constexpr float SMALL_ARMOR_WIDTH = 135;
    static constexpr float SMALL_ARMOR_HEIGHT = 55;
    static constexpr float LARGE_ARMOR_WIDTH = 230; // 230 //l 223
    static constexpr float LARGE_ARMOR_HEIGHT = 49; //45 //l 54
 
    // Four vertices of armor in 3d
    std::vector<cv::Point3f> small_armor_points_;
    std::vector<cv::Point3f> small_armor_9points_;
    std::vector<cv::Point3f> large_armor_points_;
    std::vector<cv::Point3f> large_armor_9points_;
    Eigen::Vector3d          _cam_bias;
    double                   angle  = 0.2 * CV_PI;
    float                    lastr  = 0.14;
    float                    radius = 0.24;
    cv::Point2f              ccc    = cv::Point2f(3, 3);
    std::vector<cv::Point3f> armorss;
    std::vector<cv::Point3f> armorss1;
};

}  // namespace hnurm

#endif  // ARMOR_DETECTOR_PNP_SOLVER_HPP_
                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                